#include <iostream>
#include <Eigen/Dense>
#include <opencv2/core/core.hpp>
#include <opencv2/core/eigen.hpp>
#include <ceres/ceres.h>
#include <Eigen/Core>
#include <chrono>

using namespace std;

template <typename Derived>
Eigen::Matrix<typename Derived::Scalar, 3, 3> skewSymmetric(const Eigen::MatrixBase<Derived> &q)
{
    Eigen::Matrix<typename Derived::Scalar, 3, 3> ans;
    ans << typename Derived::Scalar(0), -q(2), q(1),
    q(2), typename Derived::Scalar(0), -q(0),
    -q(1), q(0), typename Derived::Scalar(0);
    return ans;
}

template <typename Derived>
Eigen::Matrix<typename Derived::Scalar, 3, 2> Bx(const Eigen::MatrixBase<Derived> &q)
{
    Eigen::Matrix<typename Derived::Scalar, 3, 2> ans;
    ans << typename Derived::Scalar(1) - q.x() * q.x() / (typename Derived::Scalar(1) + q.z()), -q.x() * q.y() / (typename Derived::Scalar(1) + q.z()),
    -q.x() * q.y() / (typename Derived::Scalar(1) + q.z()), typename Derived::Scalar(1) - q.y() * q.y() / (typename Derived::Scalar(1) + q.z()),
    -q.x(), -q.y();
    return ans;
}

template <typename Derived>
Eigen::Quaternion<typename Derived::Scalar> deltaQ(const Eigen::MatrixBase<Derived> &theta)
{
    typedef typename Derived::Scalar Scalar_t;

    Eigen::Quaternion<Scalar_t> dq;
    Eigen::Matrix<Scalar_t, 3, 1> half_theta = theta;
    half_theta /= static_cast<Scalar_t>(2.0);
    dq.w() = static_cast<Scalar_t>(1.0);
    dq.x() = half_theta.x();
    dq.y() = half_theta.y();
    dq.z() = half_theta.z();
    return dq;
}

class PostureLocalParameterization : public ceres::LocalParameterization
{
    bool Plus(const double *x, const double *delta, double *x_plus_delta) const
    {
        Eigen::Map<const Eigen::Quaterniond> _q(x);

        Eigen::Quaterniond dq = deltaQ(Eigen::Map<const Eigen::Vector3d>(delta));

        Eigen::Map<Eigen::Quaterniond> q(x_plus_delta);

        q = (_q * dq).normalized();

        return true;
    }
    bool ComputeJacobian(const double *x, double *jacobian) const
    {
        Eigen::Map<Eigen::Matrix<double, 4, 3, Eigen::RowMajor>> j(jacobian);
        j.topRows<3>().setIdentity();
        j.bottomRows<1>().setZero();

        return true;
    }
    virtual int GlobalSize() const { return 4; };
    virtual int LocalSize() const { return 3; };
};

class AxisLocalParameterization : public ceres::LocalParameterization
{
    bool Plus(const double *x, const double *delta, double *x_plus_delta) const
    {
        // Assuming the input vectors are already normalized unit vectors
        Eigen::Map<const Eigen::Vector3d> v(x);
        Eigen::Map<const Eigen::Vector2d> derta_v(delta);

        Eigen::Vector3d temp = Bx(v) * derta_v;
        Eigen::AngleAxisd rotation_temp(temp.norm(), temp.normalized());
        Eigen::Matrix3d derta_rotation = rotation_temp.toRotationMatrix();

        Eigen::Map<Eigen::Vector3d> result(x_plus_delta);
        result = derta_rotation * v;

        return true;
    }
    bool ComputeJacobian(const double *x, double *jacobian) const
    {
        Eigen::Map<Eigen::Matrix<double, 3, 2, Eigen::RowMajor>> j(jacobian);
        j.topRows<2>().setIdentity();
        j.bottomRows<1>().setZero();

        return true;
    }
    virtual int GlobalSize() const { return 3; };
    virtual int LocalSize() const { return 2; };
};

class CalibGimbalRFactor : public ceres::SizedCostFunction<1, 4, 1, 3>
{
public:
    CalibGimbalRFactor(const Eigen::Matrix3d& observed_R_) : observed_R(observed_R_) { }

    virtual bool Evaluate(double const* const* parameters, double* residuals, double** jacobians) const
    {
        Eigen::Quaterniond q_be(parameters[0][3], parameters[0][0], parameters[0][1], parameters[0][2]);
        double theta = parameters[1][0];
        Eigen::Vector3d a(parameters[2][0], parameters[2][1], parameters[2][2]);
        for( ; !(theta >= -CV_PI && theta <= CV_PI); )
        {
            if(theta < -CV_PI)
                theta += 2 * CV_PI;
            else if(theta > CV_PI)
                theta -= 2 * CV_PI;
        }
        // cout << "theta in evaluate = " << theta << endl;

        Eigen::Matrix3d R_be_so3 = q_be.normalized().toRotationMatrix();
        Eigen::AngleAxisd rotation_ec(theta, a);
        Eigen::Matrix3d R_ec_so3 = rotation_ec.toRotationMatrix();
        Eigen::Vector3d y(1, 1, 1);
        Eigen::Matrix3d model_ = R_ec_so3 * R_be_so3;
        
        Eigen::Vector3d distance = model_ * y - observed_R * y;//ry
        residuals[0] = distance.norm();//   ||ry||2
        Eigen::Vector3d reduce = distance / distance.norm();

        if (jacobians)
        {
            if (jacobians[0])
            {
                Eigen::Map<Eigen::Matrix<double, 1, 4, Eigen::RowMajor>> J_R_be(jacobians[0]);
                J_R_be << - reduce.transpose() * R_ec_so3 * R_be_so3 * skewSymmetric(y), 0;
            }

            if (jacobians[1])
            {
                Eigen::Matrix3d temp = sin(theta) * Eigen::Matrix3d::Identity() + (1 + sin(theta)) * a * a.transpose() + cos(theta) * skewSymmetric(a);
                jacobians[1][0] = reduce.transpose() * temp * R_be_so3 * y;
            }

            if (jacobians[2])
            {
                Eigen::Map<Eigen::Matrix<double, 1, 3>> J_a(jacobians[2]);
                Eigen::Matrix3d Jr;
                if(theta != 0)
                    Jr = sin(theta) / theta * Eigen::Matrix3d::Identity() + (1 - sin(theta) / theta) * a * a.transpose() - (1 - cos(theta)) / theta * skewSymmetric(a);
                else
                    Jr = Eigen::Matrix3d::Identity();
                //cout << "Jr = " << Jr << endl;
                J_a << reduce.transpose() * R_ec_so3 * Jr * skewSymmetric(R_be_so3 * y) * skewSymmetric(theta * a) * Bx(a), 0;
            }
        }
        return true;
    }

private:
    Eigen::Matrix3d observed_R;
};

class CalibGimbalTFactor : public ceres::SizedCostFunction<1, 3, 3, 1, 3>
{
public:
    CalibGimbalTFactor(const Eigen::Vector3d& observed_T_) : observed_T(observed_T_) { }

    virtual bool Evaluate(double const* const* parameters, double* residuals, double** jacobians) const
    {
        Eigen::Vector3d T_be(parameters[0][0], parameters[0][1], parameters[0][2]);
        Eigen::Vector3d T_ec(parameters[1][0], parameters[1][1], parameters[1][2]);
        double theta = parameters[2][0];
        Eigen::Vector3d a(parameters[3][0], parameters[3][1], parameters[3][2]);
        for( ; !(theta >= -CV_PI && theta <= CV_PI); )
        {
            if(theta < -CV_PI)
                theta += 2 * CV_PI;
            else if(theta > CV_PI)
                theta -= 2 * CV_PI;
        }
        // cout << "theta in evaluate = " << theta << endl;

        Eigen::AngleAxisd rotation_ec(theta, a);
        Eigen::Matrix3d R_ec_so3 = rotation_ec.toRotationMatrix();
        Eigen::Vector3d model = T_be + R_ec_so3 * T_ec;
        
        Eigen::Vector3d distance = model - observed_T;//ry
        residuals[0] = distance.norm();//   ||ry||2
        Eigen::Vector3d reduce = distance / distance.norm();

        if (jacobians)
        {
            if (jacobians[0])
            {
                Eigen::Map<Eigen::Matrix<double, 1, 3, Eigen::RowMajor>> J_T_be(jacobians[0]);
                J_T_be = reduce.transpose();
            }

            if (jacobians[1])
            {
                Eigen::Map<Eigen::Matrix<double, 1, 3, Eigen::RowMajor>> J_T_ec(jacobians[1]);
                J_T_ec = reduce.transpose() * R_ec_so3;
            }

            if (jacobians[2])
            {
                Eigen::Matrix3d temp = -sin(theta) * Eigen::Matrix3d::Identity() + (1 + sin(theta)) * a * a.transpose() + cos(theta) * skewSymmetric(a);
                jacobians[2][0] = reduce.transpose() * temp * T_ec;
            }

            if (jacobians[3])
            {
                Eigen::Map<Eigen::Matrix<double, 1, 3>> J_a(jacobians[3]);
                Eigen::Matrix3d Jr;
                if(theta != 0)
                    Jr = sin(theta) / theta * Eigen::Matrix3d::Identity() + (1 - sin(theta) / theta) * a * a.transpose() - (1 - cos(theta)) / theta * skewSymmetric(a);
                else
                    Jr = Eigen::Matrix3d::Identity();
                //cout << "Jr = " << Jr << endl;
                J_a << reduce.transpose() * R_ec_so3 * Jr * skewSymmetric(T_ec) * skewSymmetric(theta * a) * Bx(a), 0;
            }
        }
        return true;
    }

private:
    Eigen::Vector3d observed_T;
};

class Observer : public ceres::IterationCallback
{
public:
    explicit Observer() {}

    ceres::CallbackReturnType operator()(const ceres::IterationSummary& summary)
    {
        printf("iteration: %d, cost: %e, cost_change: %e \r\n", summary.iteration, summary.cost, summary.cost_change);
        return ceres::SOLVER_CONTINUE;
    }
};

int main ( int argc, char** argv )
{
    printf("helloworld");
    // 读入待优化变量初值
    double R_be[4] = {-0.5111, -0.4845, 0.4804, 0.5227};
    double a[3] = {0., 0., 1.};
    double T_be[3] = {0.01, 0.01, 0.01}, T_ec[3] = {0.01, 0.01, 0.01};
    vector<double> theta;
    vector<Eigen::Matrix3d> observed_R; // 观测值
    vector<Eigen::Vector3d> observed_T;
    std::string config_file;
    // config_file = "../config/config.yaml";
    config_file = "../config/config1.yaml";
    cv::FileStorage fsSettings(config_file, cv::FileStorage::READ);
    if(!fsSettings.isOpened())
        std::cerr << "ERROR: Wrong path to settings" << std::endl;
    int num = fsSettings["num"];
    for(int i = 0; i < num; i++)
    {
        double theta_val = fsSettings["theta_" + to_string(i)];
        theta.push_back(theta_val);

        cv::Mat observed_R_cvmat, observed_T_cvmat;
        Eigen::Matrix3d observed_R_val;
        Eigen::Vector3d observed_T_val;
        fsSettings["observed_R_" + to_string(i)] >> observed_R_cvmat;
        fsSettings["observed_T_" + to_string(i)] >> observed_T_cvmat;
        cv::cv2eigen(observed_R_cvmat, observed_R_val);
        cv::cv2eigen(observed_T_cvmat, observed_T_val);
        observed_R.push_back(observed_R_val);
        observed_T.push_back(observed_T_val);

        // cout << "i = " << i << "  theta = " << theta_val << "  observed_T_val = \r\n" 
        // << observed_T_val << endl;
    }

    // 构建最小二乘问题
    ceres::Problem problem;
    ceres::LossFunction *loss_function = new ceres::CauchyLoss(1.0);
    ceres::LocalParameterization *posture_local_parameterization = new PostureLocalParameterization();
    problem.AddParameterBlock(R_be, 4, posture_local_parameterization);
    ceres::LocalParameterization *axis_local_parameterization = new AxisLocalParameterization();
    problem.AddParameterBlock(a, 3, axis_local_parameterization);
    problem.AddParameterBlock(T_be, 3);
    problem.AddParameterBlock(T_ec, 3);
    for (int i = 0; i < num; i++)
    {
        // cout << "i = " << i <<endl;
        problem.AddParameterBlock(&theta[i], 1);
        // problem.SetParameterLowerBound(&theta[i], 0, theta[i] - 5 * 2 * CV_PI / 4096);
        // problem.SetParameterUpperBound(&theta[i], 0, theta[i] + 5 * 2 * CV_PI / 4096);
        CalibGimbalRFactor *r_factor = new CalibGimbalRFactor(observed_R[i]);
        problem.AddResidualBlock(r_factor, loss_function, R_be, &theta[i], a);
        CalibGimbalTFactor *t_factor = new CalibGimbalTFactor(observed_T[i]);
        problem.AddResidualBlock(t_factor, loss_function, T_be, T_ec, &theta[i], a);
    }

    // 配置求解器
    ceres::Solver::Options options;                 // 这里有很多配置项可以填
    options.linear_solver_type = ceres::DENSE_QR;   // 增量方程如何求解
    // options.minimizer_progress_to_stdout = true;    // 输出到cout
    options.function_tolerance = 1e-4;
    options.max_num_iterations = 100;
    options.gradient_tolerance = 1e-9;
    options.linear_solver_type = ceres::DENSE_NORMAL_CHOLESKY;
    options.use_nonmonotonic_steps = true;
    // options.dynamic_sparsity = true;
    options.callbacks.push_back(new Observer());

    ceres::Solver::Summary summary;                 // 优化信息
    chrono::steady_clock::time_point t1 = chrono::steady_clock::now();
    ceres::Solve ( options, &problem, &summary );   // 开始优化
    chrono::steady_clock::time_point t2 = chrono::steady_clock::now();
    chrono::duration<double> time_used = chrono::duration_cast<chrono::duration<double>>( t2-t1 );
    cout<<"solve time cost = "<<time_used.count()<<" seconds. "<<endl;

    // 输出结果
    cout << summary.BriefReport() <<endl;
    // cout << summary.FullReport() <<endl;
    cout << "x y z w = ";
    for(auto R_be_ : R_be)
        cout << R_be_ << " ";
    cout << endl;
    cout << "theta = ";
    for(auto theta_ : theta)
        cout << theta_ << " ";
    cout << endl;
    cout << "a = ";
    for(auto a_ : a)
        cout << a_ << " ";
    cout << endl;
    cout << "T_be = ";
    for(auto T_be_ : T_be)
        cout << T_be_ << " ";
    cout << endl;
    cout << "T_ec = ";
    for(auto T_ec_ : T_ec)
        cout << T_ec_ << " ";
    cout << endl;

    return 0;
}




